from typing import TYPE_CHECKING, Union

import torch
from zkl_metrics import MeanMetric, Metric
from zkl_training import ProcessingTaskPlugin

if TYPE_CHECKING:
    from .training import GPTTraining, Output
    from .validating import GPTValidating


class MetricsPlugin(ProcessingTaskPlugin):
    def __init__(self):
        super().__init__()
        self.ce_metric = MeanMetric()
        self.acc_metric = MeanMetric()

    @property
    def task(self) -> Union['GPTTraining', 'GPTValidating']:
        task = super().task
        from .training import GPTTraining
        from .validating import GPTValidating
        assert isinstance(task, (GPTTraining, GPTValidating))
        return task

    # task

    def on_after_process(self, output: 'Output'):
        batch_tokens_ce, batch_tokens_acc, batch_tokens_mask, _ = output
        # [..., context_size]

        ce = torch.masked.sum(batch_tokens_ce, dtype=torch.float64, mask=batch_tokens_mask)
        acc = torch.masked.sum(batch_tokens_acc, dtype=torch.float64, mask=batch_tokens_mask)
        weight = torch.sum(batch_tokens_mask, dtype=torch.float64)
        # []

        self.ce_metric.update(ce, weight)
        self.acc_metric.update(acc, weight)

    # metrics

    @property
    def metrics(self) -> dict[str, Metric]:
        return {
            attr_name[:-len('_metric')]: attr_value
            for attr_name, attr_value in vars(self).items()
            if attr_name.endswith('_metric') and isinstance(attr_value, Metric)}

    def collect_metrics_value(self, reset: bool = True) -> dict[str, float]:
        metrics_value = {}
        for metric_name, metric in self.metrics.items():
            metrics_value[metric_name] = metric.compute().cpu().item()
            if reset:
                metric.reset()
        return metrics_value
